#!/usr/bin/env python
#-*- coding: UTF-8 -*-
# manager.py: task manager core implementaion
#
# Copyright (C) 2008-2010 Hily Jiang <hilyjiang@gmail.com>
#
# This file is part of dtask.
#
# This software is provided 'as-is', without any express or implied
# warranty.  In no event will the authors be held liable for any damages
# arising from the use of this software.
#
# Permission is granted to anyone to use this software for any purpose,
# including commercial applications, and to alter it and redistribute it
# freely, subject to the following restrictions:
#
# 1. The origin of this software must not be misrepresented; you must not
#    claim that you wrote the original software. If you use this software
#    in a product, an acknowledgment in the product documentation would be
#    appreciated but is not required.
# 2. Altered source versions must be plainly marked as such, and must not be
#    misrepresented as being the original software.
# 3. This notice may not be removed or altered from any source distribution.

"""
Task Manager core implementaion.
"""

import sys
import threading
import sqlite3
from time import time


class TaskManager(object):
    """Task Manager, it hold a task buffer."""

    def __init__(self, dbpath=':memory:', taskTimerInterval=5,
                 resetIntervals=(60,)*5, clearIntervals=(86400,)*3,
                 maxRetries=5, encoding='utf-8'):
        """
        dbpath:
            sqlite3 database save path, default ':memory:' means run in memory.
        taskTimerInterval:
            interval of the timer which check if the task is timeout
        resetInterval:
            a interval tuple to reset the unfinished task.
            the items in the tuple means the intervals to reset tasks:
               (scheduling, scheduled, processing, processed, transfering)
            if the task timeout is specified, this will not effect.
        clearInterval:
            a interval tuple to clear tasks.
            the items in the tuple mean the intervals to clean tasks:
               (finished, returned, failed)
        maxRetries:
            if a task retried over max retry times, it will be signed as failed.
        encoding:
            default encoding of sqlite3 output is unicode, we always use utf-8.
        """
        self.dbpath = dbpath
        self.taskTimerInterval = taskTimerInterval
        self.resetTypes = ('scheduling', 'scheduled', 'processing',
                           'processed', 'transfering')
        self.clearTypes = ('fininshed', 'returned', 'failed')
        self.resetIntervals = dict(zip(self.resetTypes, resetIntervals))
        self.clearIntervals = dict(zip(self.clearTypes, clearIntervals))
        self.maxRetries = maxRetries
        self.encoding = encoding

        # timers
        self.taskTimer = self.__startTaskTimer()
        self.resetTimers = {}
        self.clearTimers = {}
        for timerType in self.resetTypes:
            self.__startResetTimer(timerType)
        for timerType in self.clearTypes:
            self.__startClearTimer(timerType)

        # blocking implemetation resource
        self.resultEvents = {}
        self.tasksem = threading.Semaphore(0)

        # database operation implementation resource
        self.dbThreadLock = threading.Semaphore(0)
        self.dbThreadLock.release()
        self.dbOpEvent = threading.Event()
        self.dbReturnEvent = threading.Event()
        self.dbOpFunc = None
        self.dbOpArgs = None
        self.dbOpKwrgs = None
        self.dbReturn = None

        # Because sqlite3 does not support cross-thread connection sharing,
        # so we create an individual thread to handle sqlite3 operations.
        self.dbthread = threading.Thread(target=self.__dbThreadFunc)
        self.dbthread.setDaemon(True)
        self.dbthread.start()

    def __dbThreadFunc(self):
        """Thread function for sqlite3 operations."""
        self.conn = sqlite3.connect(self.dbpath)
        self.cur = self.conn.cursor()
        try:
            # create table for task manager
            # taskid:     task ID
            # tasktype:   task type
            # params:     task params
            # status:     task status
            #     'new': task is newly created
            #     'scheduling': task is scheduling by scheduler
            #     'scheduled':  task has scheduled by scheduler
            #     'processing': task is processing by some node
            #     'processed':  task has processed by some node
            #     'transfering': result is being transfered
            #     'finished': result has write to tasks table
            #     'returned': result has return to application
            #     'failed': task that retries over the limit or timeout
            # retries:  retry times
            # appname:  application name
            # appip:    application ip
            # appport:  application port
            # appast:   application assign time
            # apprtt:   application return time
            # schname:  scheduler name
            # schip:    scheduler ip
            # schport:  scheduler port
            # schast:   scheduler assign time
            # schrtt:   scheduler return time
            # nodmip:   node manager ip
            # nodmport: node manager port
            # nodmast:  node manager assign time
            # nodmrtt:  node manager return time
            # nodname:  node name
            # nodip:    node ip
            # nodport:  node port
            # nodast:   node assign time
            # nodrtt:   node return time
            # result:   task result
            self.cur.execute('''
                CREATE TABLE tasks(
                    taskid INTEGER PRIMARY KEY AUTOINCREMENT,
                    tasktype CHAR(20),
                    timeout INTEGER,
                    params CHAR(200),
                    status CHAR(20),
                    retries INTEGER,
                    appname CHAR(20),
                    appip CHAR(20),
                    appport INTEGER,
                    appast INTEGER,
                    apprtt INTEGER,
                    schname CHAR(20),
                    schip CHAR(20),
                    schport INTEGER,
                    schast INTEGER,
                    schrtt INTEGER,
                    nodmip CHAR(20),
                    nodmport INTEGER,
                    nodmast INTEGER,
                    nodmrtt INTEGER,
                    nodname CHAR(20),
                    nodip CHAR(20),
                    nodport INTEGER,
                    nodast INTEGER,
                    nodrtt INTEGER,
                    result TEXT
                )
            ''')
            self.conn.commit()
        except: # table exists
            # reset task status to 'new' if it's status is not 'finished'
            # or 'returned'
            self.cur.execute('UPDATE tasks SET status = "new"' \
                             ' WHERE status<>"finished"' \
                             ' AND status<>"returned"' \
                             ' AND status<>"failed"')
            self.conn.commit()

            # create task semaphores for the task that haven't return result
            self.cur.execute('SELECT taskid FROM tasks' \
                             ' WHERE status<>"returned"')
            taskIds = self.cur.fetchall()
            for taskId in taskIds:
                taskId = taskId[0]
                self.resultEvents[taskId] = threading.Event()

            # add semaphore for 'new' status task
            self.cur.execute('SELECT COUNT(1) FROM tasks' \
                             ' WHERE status=="new"')
            count = self.cur.fetchone()[0]
            self.tasksem = threading.Semaphore(count)

        while True:
            self.dbOpEvent.wait()
            self.dbOpEvent.clear()
            self.dbReturn = self.dbOpFunc(self, *self.dbOpArgs,
                                          **self.dbOpKwargs)
            self.dbReturnEvent.set()

    def __dbThread(func):
        """Decoration function to wrap functions that operate sqlite3.
           The functions decorated will then executed in the db thread."""
        def wrapFunc(self, *args, **kwargs):
            self.dbThreadLock.acquire()
            self.dbOpFunc = func
            self.dbOpArgs = args
            self.dbOpKwargs = kwargs
            self.dbOpEvent.set()
            self.dbReturnEvent.wait()
            rt = self.dbReturn
            self.dbReturnEvent.clear()
            self.dbThreadLock.release()
            return rt
        return wrapFunc

    def __startTaskTimer(self):
        """Start a timer to check timeout of tasks ."""
        timer = threading.Timer(self.taskTimerInterval, self.__taskTimerFunc)
        timer.setDaemon(True)
        timer.start()

    def __startResetTimer(self, timerType):
        """Start a timer to reset specified type of tasks."""
        timer = threading.Timer(self.resetIntervals[timerType],
                              self.__resetTimerFunc, args=(timerType,))
        timer.setDaemon(True)
        timer.start()
        self.resetTimers[timerType] = timer

    def __startClearTimer(self, timerType):
        """Start a timer to clear specified type of tasks."""
        timer = threading.Timer(self.clearIntervals[timerType],
                              self.__clearTimerFunc, args=(timerType,))
        timer.setDaemon(True)
        timer.start()
        self.clearTimers[timerType] = timer

    def __taskTimerFunc(self):
        """Check task timeout function. It is intervally executed."""
        self.__checkTaskTimeout()
        self.__startTaskTimer()

    def __resetTimerFunc(self, timerType):
        """Reset task function. It is intervally executed."""
        self.__resetTask(timerType)
        self.__startResetTimer(timerType)

    def __clearTimerFunc(self, timerType):
        """Clear task function. It is intervally executed."""
        self.__clearTask(timerType)
        self.__startClearTimer(timerType)

    @__dbThread
    def __checkTaskTimeout(self):
        """Check if task has timeout. Set status to failed if timeout."""
        nowTime = int(time())

        sql = 'SELECT COUNT(1) FROM tasks WHERE ?>schast+timeout' \
              ' AND status="new"'
        self.cur.execute(sql, (nowTime,))
        failCount = self.cur.fetchone()[0]

        # update task semaphore
        for i in range(failCount):
                self.tasksem.acquire()

        sql = 'SELECT taskid FROM tasks WHERE ?>schast+timeout' \
              ' AND status<>"failed"'
        self.cur.execute(sql, (nowTime,))
        taskIds = self.cur.fetchall()
        sql = 'UPDATE tasks SET status="failed"' \
              ' WHERE ?>schast+timeout'
        self.cur.execute(sql, (nowTime,))
        self.conn.commit()

        # return result of those failed task
        for line in taskIds:
            taskId = line[0]
            if self.resultEvents.has_key(taskId):
                self.resultEvents[taskId].set()

    @__dbThread
    def __resetTask(self, status):
        """Reset specified status of tasks."""
        sql = 'UPDATE tasks SET status="new", retries=retries+1,' \
              ' schname="", schip="", schport=0, schast=0, schrtt=0,' \
              ' nodmip="", nodmport="", nodmast=0, nodmrtt=0,' \
              ' nodname="", nodip="", nodport="", nodast=0, nodrtt=0' \
              ' WHERE status=? AND timeout=0'
        self.cur.execute(sql, (status,))
        self.conn.commit()
        resetCount = self.cur.rowcount

        # set status to failed if retry times over maxRetries
        if resetCount > 0:
            sql = 'SELECT taskid FROM tasks WHERE retries>?' \
                  ' AND status<>"failed"'
            self.cur.execute(sql, (self.maxRetries,))
            taskIds = self.cur.fetchall()
            sql = 'UPDATE tasks SET status="failed" ' \
                  ' WHERE retries>?'
            self.cur.execute(sql, (self.maxRetries,))
            self.conn.commit()
 
            # return result of those failed task
            for line in taskIds:
                taskId = line[0]
                if self.resultEvents.has_key(taskId):
                    self.resultEvents[taskId].set()

            # new task avaliable
            resetCount -= self.cur.rowcount # minus failed counts
            for i in range(resetCount):
                self.tasksem.release()

    @__dbThread
    def __clearTask(self, status):
        """Clear specified status of tasks."""
        self.cur.execute('DELETE FROM tasks WHERE status=?', (status,))
        self.conn.commit()

    @__dbThread
    def newTask(self, appIp, appPort, taskType, taskTimeout, taskParams, appName):
        """Add a new task to task manager. Return an ID identifies the task.
           Interface for application."""
        appPort = int(appPort)
        taskTimeout = int(taskTimeout)
        sql = 'INSERT INTO tasks VALUES(NULL, ?, ?, ?, "new", 0,' \
              ' ?, ?, ?, ?, 0, "", "", 0, 0, 0, "", 0, 0, 0,' \
              ' "", "", 0, 0, 0, "")'
        self.cur.execute(sql, (taskType, taskTimeout, taskParams,
                         appName, appIp, appPort, int(time())))
        self.conn.commit()
        rowid = self.cur.lastrowid
        sql = 'SELECT taskid FROM tasks WHERE ROWID=?'
        taskId = self.cur.execute(sql, (rowid,)).fetchone()[0]

        self.resultEvents[taskId] = threading.Event()
        self.tasksem.release()

        # return row id as task id
        return taskId

    def __encodeResult(self, result):
        """Encode task result."""
        needEncode = ('tasktype', 'params', 'result',
                      'appname', 'schname', 'nodname')
        for key in needEncode:
            if result.has_key(key):
                result[key].encode(self.encoding)
        return result

    @__dbThread
    def __getResultFromDB(self, taskId, fullInfo, appIp, appPort):
        """Get task result from database."""
        taskId = int(taskId)
        appPort = int(appPort)
        if fullInfo:
            fields = ('tasktype', 'timeout', 'params', 'status', 'retries',
                      'appname', 'appip', 'appport', 'appast', 'apprtt',
                      'schname', 'schip', 'schport', 'schast', 'schrtt',
                      'nodmip', 'nodmport', 'nodmast', 'nodmrtt',
                      'nodname', 'nodip', 'nodport', 'nodast', 'nodrtt',
                      'result')
            fieldsStr = ','.join(fields)
        else:
            fields = ('result', )
            fieldsStr = 'result'
        self.cur.execute('SELECT %s FROM tasks WHERE taskid=?'\
                         ' AND appip=? AND appport=?' % fieldsStr,
                         (taskId, appIp, appPort))
        result = self.cur.fetchone()
        if result != None:
            result = dict(zip(fields, result))
            result = self.__encodeResult(result)
        del self.resultEvents[taskId]
        self.cur.execute('UPDATE tasks SET status="returned",' \
                         ' apprtt=? WHERE taskid=?', (int(time()), taskId))
        self.conn.commit()
        return result

    def getResult(self, taskId, fullInfo, appIp, appPort):
        """Get task result. If task haven't finished it may block."""
        taskId = int(taskId)
        if self.resultEvents.has_key(taskId):
            self.resultEvents[taskId].wait()
            self.resultEvents[taskId].clear()
            return self.__getResultFromDB(taskId, fullInfo, appIp, appPort)
        else:
            return ''

    @__dbThread
    def writeResult(self, taskId, schIp, schPort, result):
        """Write result to database, and unblock the corresponding semaphore."""
        if schIp=='': return
        taskId = int(taskId)
        schPort = int(schPort)
        sql = 'UPDATE tasks SET result=?, status="finished",' \
              ' schrtt=? WHERE taskId=? AND schip=? AND schport=?'
        self.cur.execute(sql, (result, int(time()), taskId, schIp, schPort))
        self.conn.commit()
        if self.resultEvents.has_key(taskId):
            self.resultEvents[taskId].set()

    @__dbThread
    def __getTaskFromDB(self, schIp, schPort, schName):
        """Get a task from database."""
        schPort = int(schPort)
        self.cur.execute('SELECT taskid, tasktype, params FROM tasks' \
                         ' WHERE status="new" ORDER BY appast')
        taskInfo = self.cur.fetchone() # (taskId, taskType, taskParams)
        taskId, taskType, taskParams = taskInfo
        taskType = taskType.encode(self.encoding)
        taskParams = taskParams.encode(self.encoding)
        if taskInfo:
            sql = 'UPDATE tasks SET schip=?, schport=?,' \
                  ' schname=?, schast=?,' \
                  ' status="scheduling" WHERE taskid=?'
            self.cur.execute(sql, (schIp, int(schPort), schName,
                                   int(time()), taskId))
            self.conn.commit()
        return (taskId, taskType, taskParams)

    def getTask(self, schIp, schPort, schName):
        """Get a task from the task list.
           Interface for scheduler."""
        schPort = int(schPort)
        self.tasksem.acquire()
        return self.__getTaskFromDB(schIp, schPort, schName)

    @__dbThread
    def updateNodeManagerStatus(self, taskId, schIp, schPort, nodmIp=None,
                                nodmPort=None, nodmAst=None, nodmRtt=None):
        """Update task status when task is transfered between node manager
           and scheduler."""
        taskId = int(taskId)
        schPort = int(schPort)
        nodmPort = int(nodmPort)
        nodmAst = int(nodmAst)
        nodmRtt = int(nodmRtt)
        if nodmRtt:
            status = 'transfering'
        elif nodmAst:
            status = 'scheduled'
        setInfo = 'status="%s"' % status
        try:
            if nodmIp and nodmPort and nodmAst:
                setInfo += ', nodmip="%s", nodmport=%d, nodmast=%d' % \
                           (nodmIp, nodmPort, nodmAst)
            if nodmRtt:
                setInfo += ', nodmrtt=%d' % nodmRtt
        except Exception, e:
            print '[MGR]', sys._getframe().f_lineno, e
        sql = 'UPDATE tasks SET %s WHERE taskid=?' \
              ' AND schip=? AND schport=?' % setInfo
        self.cur.execute(sql, (taskId, schIp, schPort))
        self.conn.commit()

    @__dbThread
    def updateNodeStatus(self, taskId, schIp, schPort, nodName=None,
                         nodIp=None, nodPort=None, nodAst=None, nodRtt=None):
        """Update task status when task is transfered between node manager
           and node."""
        taskId = int(taskId)
        schPort = int(schPort)
        nodPort = int(nodPort)
        nodAst = int(nodAst)
        nodRtt = int(nodRtt)
        if nodRtt:
            status = 'processed'
        elif nodAst:
            status = 'processing'
        setInfo = 'status="%s"' % status
        try:
            if nodIp and nodPort and nodName and nodAst:
                setInfo += ', nodip="%s", nodport=%d, nodname="%s",' \
                           ' nodast=%d' % (nodIp, nodPort, nodName, nodAst)
            if nodRtt:
                setInfo += ', nodrtt=%d' % nodRtt
        except Exception, e:
            print '[MGR]', sys._getframe().f_lineno, e
        sql = 'UPDATE tasks SET %s WHERE taskid=?' \
              ' AND schip=? AND schport=?' % setInfo
        self.cur.execute(sql, (taskId, schIp, schPort))
        self.conn.commit()
